#!/usr/bin/env python

import sys;
import adodb;

class CleanUp:

	def __init__(self, dbname, username):
		self.conn = adodb.NewADOConnection('postgres')
		self.conn.Connect('user=%s dbname=%s' % (username, dbname))

	def doit(self):

		self.cursor = self.conn.Execute('SELECT DISTINCT artist, album, track, mbid from Scrobbles WHERE stid is NULL');

		while not self.cursor.EOF:
			row = self.cursor.GetRowAssoc(0)
			artist = row['artist']
			album = row['album']
			track = row['track']
			mbid = row['mbid']
			stid = self.lookupStid(artist,album,track,mbid)
			if not stid:
				tid = self.lookupTrack(artist, album, track, mbid)
				if tid:
					print "Adding stid for %s/%s/%s/%s/%s" % (tid,artist,album,track,mbid)
					self.acursor = self.conn.Execute("INSERT INTO Scrobble_Track (artist,album,name,mbid,track) VALUES(lower(%(artist)s),lower(%(album)s),lower(%(name)s),lower(%(mbid)s),%(track)s)", {'artist': artist, 'album': album, 'name': track, 'mbid': mbid, 'track': tid})
				else:
					print "Adding tid for %s/%s/%s/%s" % (artist,album,track,mbid)
					self.acursor = self.conn.Execute("INSERT INTO Track (artist,album,name,mbid) VALUES(%(artist)s,%(album)s,%(name)s,%(mbid)s)", {'artist': artist, 'album': album, 'name': track, 'mbid': mbid})
			else:
				print "Setting %d to match %s/%s/%s/%s" % (stid,artist,album,track,mbid)
				self.scursor = self.conn.Execute("UPDATE Scrobbles SET stid = %%(stid)s WHERE track = %%(track)s AND artist = %%(artist)s AND %s AND %s AND stid IS NULL" % (('lower(album) = lower(%(album)s)' if album else 'album IS NULL'), ('lower(mbid) = lower(%(mbid)s)' if mbid else 'mbid IS NULL')), {'stid': stid, 'track': track, 'artist': artist, 'album': album, 'mbid': mbid})
			self.cursor.MoveNext()

	def lookupStid(self, artist, album, track, mbid):
		self.lucursor = self.conn.Execute('SELECT * from Scrobble_Track WHERE lower(artist) = lower(%%(artist)s) AND %s AND lower(name) = lower(%%(track)s) AND %s' % (('lower(album) = lower(%(album)s)' if album else 'album IS NULL'), ('lower(mbid) = lower(%(mbid)s)' if mbid else 'mbid IS NULL')), {'artist': artist, 'album': album, 'track': track, 'mbid': mbid});
		if not self.lucursor.EOF:
			lrow = self.lucursor.GetRowAssoc(0)
			return lrow['id']
		else:
			return None

	def lookupTrack(self, artist, album, track, mbid):
		self.tucursor = self.conn.Execute('SELECT * from Track WHERE lower(artist) = lower(%%(artist)s) AND %s AND lower(name) = lower(%%(track)s) AND %s' % (('lower(album) = lower(%(album)s)' if album else 'album IS NULL'), ('lower(mbid) = lower(%(mbid)s)' if mbid else 'mbid IS NULL')), {'artist': artist, 'album': album, 'track': track, 'mbid': mbid});
		if not self.tucursor.EOF:
			trow = self.tucursor.GetRowAssoc(0)
			return trow['id']
		else:
			return None

	def close(self):
		self.cursor.Close()
		self.conn.Close() 

if __name__ == "__main__":

	if len(sys.argv) != 3:
		print "Usage: setstids <pgsql dbname> <pgsql username>"
		sys.exit(1)
	
	cleaner = CleanUp(sys.argv[1], sys.argv[2])
	cleaner.doit()
	cleaner.close()
